import numpy as np
from scipy.stats import f, ncf

def compute_variance_of_means(means):
    """
    按照 nQuery 界面给出的公式:
    V = (1/M) * Σ((μ_i - μ_bar)^2),
    其中 M = len(means), μ_bar = 均值的平均值
    """
    means = np.array(means, dtype=float)
    M = len(means)
    mu_bar = np.mean(means)
    sum_of_squares = np.sum((means - mu_bar)**2)
    V = sum_of_squares / M
    return V

def repeated_measures_sample_size(
    alpha: float,
    M: int,
    V: float,
    sd: float,
    cor: float,
    power: float,
    step: float = 1.0,
    max_n: int = 1000
):
    """
    基于单因素重复测量 ANOVA 的简易样本量估计函数:
      - alpha: 显著性水平 (如 0.05)
      - M: 重复测量水平数 (如有4个时间点, M=4)
      - V: "Variance of means" (从 compute_variance_of_means 计算或手动指定)
      - sd: 每水平的标准差
      - cor: 水平间的相关系数 (rho)
      - power: 目标检验功效(如 0.90 表示 90%)
      - step: 在迭代中每次增加的 n 步幅
      - max_n: 最大循环搜索范围, 防止无限增大

    其中, 效应量 es = V / [sd^2 * (1 - cor)]
    非中心参数 ncp = n * M * es
    自由度 df1 = M - 1, df2 = (M - 1)*(n - 1)
    功效 = 1 - F_{ncf}(f临界值; df1, df2, ncp)
    """
    # 基本检查
    if not (0 < alpha < 1):
        raise ValueError("alpha 必须在 (0,1) 范围内.")
    if M < 2:
        raise ValueError("M(重复测量水平数) 不能 < 2.")
    if sd <= 0:
        raise ValueError("sd 必须 > 0.")
    if not (-1 <= cor <= 1):
        raise ValueError("cor 必须在 [-1,1] 范围内.")
    if not (0 < power < 1):
        raise ValueError("power 必须在 (0,1) 范围内, 如 0.90.")

    # 计算效应量 es
    es = V / (sd**2 * (1 - cor))

    # 开始搜索满足 power 的最小 n
    best_n = None
    for n in np.arange(2, max_n + step, step):
        # 计算自由度
        df1 = M - 1
        df2 = (M - 1) * (n - 1)
        if df2 <= 0:
            continue

        # 非中心参数
        ncp = n * M * es

        # 中心F分布下 1 - alpha 分位数
        f_crit = f.ppf(1 - alpha, df1, df2)
        if np.isinf(f_crit) or np.isnan(f_crit):
            continue

        # 在非中心F分布下, 统计量>f_crit 的概率 => 功效
        # 故 pw = 1 - ncf.cdf(f_crit, df1, df2, ncp)
        cdf_val = ncf.cdf(f_crit, df1, df2, ncp)
        pw = 1 - cdf_val

        if pw >= power:
            best_n = int(np.ceil(n))
            break

    return best_n


# ---------------------- 测试示例 ----------------------
if __name__ == "__main__":
    # 1) 先计算 V (Variance of means)
    means_example = [55, 56.5, 58, 59.5]
    V_example = compute_variance_of_means(means_example)
    print(f"根据公式 V = Σ(μ_i - μ_bar)^2 / M 得到: V = {V_example:.4f}")

    # 2) 设定其他参数: alpha, sd, cor, power
    alpha_test = 0.05  # 显著性水平
    M_test = 4         # 4 个水平 (对应 4 个时间点)
    sd_test = 10       # 每个水平的标准差
    cor_test = 0.7     # 水平间相关系数
    power_test = 0.90  # 90% 功效

    # 3) 调用重复测量样本量估计函数
    n_est = repeated_measures_sample_size(
        alpha=alpha_test,
        M=M_test,
        V=V_example,
        sd=sd_test,
        cor=cor_test,
        power=power_test,
        step=1.0,
        max_n=200
    )

    print(f"基于上述参数估算的样本量 n = {n_est}")